#include "SharedFileName.h"
#include "Utility.h"
#include <iostream>
#include <sstream>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
#include <omp.h>
#include <unordered_map>
#include <unordered_set>
#include <set>
#include <queue>
#define MAX_DIST 10000
using namespace std;

// compute the distance between different partitions
void EstimatePartitionDistance(
	vector<vector<pair<int,int> > >& adj_list,
	unordered_map<int,int>& node_partition_map, 
	int num_partition, 
	vector<vector<int> >& indirect_dist // results has been size to num_partition
	){
	
	vector<vector<int> > direct_dist(num_partition);

	for(int i = 0; i < num_partition; i++) direct_dist[i].resize(num_partition,MAX_DIST);

	// compute direct shorest path
	for(int i = 0; i < adj_list.size(); i++){
		int x_map = node_partition_map[i];
		for(int j = 0; j < adj_list[i].size(); j++){
			int y_map = node_partition_map[adj_list[i][j].first];
			if(direct_dist[x_map][y_map] > adj_list[i][j].second){
				direct_dist[x_map][y_map] = adj_list[i][j].second;
			}
		}
	}

	vector<vector<int> > partition_direct_neighours(num_partition);
	for(int i = 0; i < num_partition; i++){
		for(int j = 0; j < num_partition; j++){
			if(i != j && direct_dist[i][j] != MAX_DIST){
				partition_direct_neighours[i].push_back(j);
			}
		}
	}

	// compute indirect path
	#pragma omp parallel for
	for(int i = 0; i < num_partition; i++){
		indirect_dist[i].resize(num_partition,MAX_DIST);

		unordered_set<int> visited_partition;
		set<pair<int,int> > candidate_queue;
		candidate_queue.insert(make_pair(0,i)); // init with 0 first, set later

		int computed_partitions = 0;
		while(computed_partitions != num_partition){
			pair<int,int> best_pair = *candidate_queue.begin();
			int best_dist = best_pair.first;
			int best_partition = best_pair.second;
			candidate_queue.erase(candidate_queue.begin());
			
			if(visited_partition.find(best_partition) != visited_partition.end()) continue;
			computed_partitions++;
			visited_partition.insert(best_partition);


			// determine the best distance
			indirect_dist[i][best_partition] = best_dist;

			for(int j = 0; j < partition_direct_neighours[best_partition].size(); j++){
				int v = partition_direct_neighours[best_partition][j];
				if(visited_partition.find(v) != visited_partition.end()) continue;

				int dist = best_dist + direct_dist[best_partition][v];
				if(dist < indirect_dist[i][v]){
					indirect_dist[i][v] = dist;
					candidate_queue.insert(make_pair(dist,v));
				}
			}			
		}

		indirect_dist[i][i] = direct_dist[i][i];
		for(int j = 0; j < num_partition; j++){
			if(i != j && indirect_dist[i][i] > 2*indirect_dist[i][j]){
				indirect_dist[i][i] = 2*indirect_dist[i][j];
			}
		}
	}

	vector<pair<double,int> > centrality;
	for (int i = 0; i < num_partition; i++)
	{
		double sum = 0;
		for (int j = 0; j < num_partition; j++)
		{
			sum += indirect_dist[i][j];
		}
		centrality.push_back(make_pair(sum,i));
	}
	sort(centrality.begin(),centrality.end(),pair_greater<double,int>());
}

void SwapNearestPartition(int partition_idx, int stride_size, vector<int>& partition_order, vector<vector<int> >& part_dist){
	// find the nearest partition
	int neighbour_partition_idx = -1;
	double neighbour_distance = MAX_DIST;
	for(int idx = partition_idx + stride_size; idx < partition_order.size(); idx += stride_size)
	{
		// compute distance between two partitions
		double average_distance = 0;
		for(int x = partition_idx; x < partition_idx + stride_size; x++)
		{
			for (int y = idx; y  < idx + stride_size; y++)
			{
				int part_x = partition_order[x];
				int part_y = partition_order[y];
				average_distance += part_dist[part_x][part_y];
			}
		}
		average_distance /= (stride_size * stride_size);
		if(average_distance < neighbour_distance){
			neighbour_partition_idx = idx;
			neighbour_distance = average_distance;
		}
	}

	// swap the partition
	for (int i = 0; i < stride_size; i++)
	{
		std::swap(partition_order[partition_idx + stride_size + i],partition_order[neighbour_partition_idx + i]);
	}
}

void TreePartition(vector<vector<pair<int,int> > >& adj_list,unordered_map<int,int>& node_partition_map,string out_dir, int num_partition){
	cout << "starting for tree partition" << endl;
	vector<vector<int> > indirect_dist(num_partition);
	EstimatePartitionDistance(adj_list,node_partition_map,num_partition,indirect_dist);

	// compute number of partition levels
	int num_level = 0;
	int temp_partition_count = num_partition;
	while(temp_partition_count > 0)
	{
		temp_partition_count /= 2;
		num_level++;
	}

	vector<int> partition_order(num_partition);
	for (int i = 0; i < partition_order.size(); i++) partition_order[i] = i; 

	// sort the partitions by building the tree
	for(int level = 0; level < num_level - 1; level++)
	{
		int stride = (1 << level);
			
		for(int partition_idx = 0; partition_idx < num_partition; partition_idx += (2*stride) )
		{
			SwapNearestPartition(partition_idx,stride,partition_order,indirect_dist);
		}
	}

	unordered_map<int,int> part_remap;
	for (int i = 0; i < num_partition; i++)
	{
		part_remap[partition_order[i]] = i;
	}
	if(part_remap.size() != num_partition)
	{
		cout << "something wrong in the partition remap" << endl;
		exit(1);
	}

	string tree_node_map_str(out_dir); tree_node_map_str.append(sample_tree_partition_node_map);
	ofstream treeNodeMapFile(tree_node_map_str.c_str(),ios::out);
	for (int i = 0; i < adj_list.size(); i++)
	{
		int part_idx = node_partition_map[i];
		treeNodeMapFile << part_remap[part_idx] << endl;
	}
	treeNodeMapFile.close();

	string part_dist_str(out_dir); part_dist_str.append(sample_tree_partition_dist);
	ofstream partDistFile(part_dist_str.c_str(),ios::out);
	for (int i = 0; i < num_partition; i++)
	{
		int x_part = partition_order[i];
		for (int j = 0; j < num_partition; j++)
		{
			int y_part = partition_order[j];
			partDistFile << indirect_dist[x_part][y_part] << " ";
		}
		partDistFile << endl;
	}
	partDistFile.close();
}

// the computation for non-hierarchical partition scheme
void NormalPartition(vector<vector<pair<int,int> > >& adj_list,unordered_map<int,int>& node_partition_map, string out_dir, int num_partition){

	vector<vector<int> > indirect_dist(num_partition);
	EstimatePartitionDistance(adj_list,node_partition_map,num_partition,indirect_dist);

	// select the user with max degree for each partition
	vector<int> user_max_degree_per_partition(num_partition);
	vector<int> max_degree_per_partition(num_partition,0);

	for(auto it = node_partition_map.begin(); it != node_partition_map.end(); it++){
		int user_id = it->first;
		int partition_id = it->second;
		int user_degree_in_partition = 0;
		for(int i = 0; i < adj_list[user_id].size(); i++){
			if(node_partition_map[adj_list[user_id][i].first] == partition_id){
				user_degree_in_partition++;
			}
		}

		if(user_degree_in_partition > max_degree_per_partition[partition_id]){
			max_degree_per_partition[partition_id] = user_degree_in_partition;
			user_max_degree_per_partition[partition_id] = user_id;
		}
	}
	for(int i = 0 ;i < num_partition; i++){
		cout << "partition " << i << " select user: " << user_max_degree_per_partition[i] << " with degree: " << max_degree_per_partition[i] << endl;
	}

	//find shorest distance to all points for the pivot nodes
	vector<vector<int> > pivot_to_all_dist(num_partition);
	int num_nodes = adj_list.size();

	#pragma omp parallel for
	for(int part_idx = 0; part_idx < num_partition; part_idx++){
		
		int source_node = user_max_degree_per_partition[part_idx];
		pivot_to_all_dist[part_idx].resize(num_nodes,MAX_DIST);
		pivot_to_all_dist[part_idx][source_node] = 0;

		set<pair<int,int> > candidate;
		for(int k = 0; k < num_nodes; k++) candidate.insert(make_pair(pivot_to_all_dist[part_idx][k],k));

		// Dijkstra algorithm
		while(!candidate.empty()){
			int pivot_node = candidate.begin()->second;
			candidate.erase(candidate.begin());

			for(int i = 0; i < adj_list[pivot_node].size(); i++){
				int v = adj_list[pivot_node][i].first;
				int alt = pivot_to_all_dist[part_idx][pivot_node] + adj_list[pivot_node][i].second;
				if(alt < pivot_to_all_dist[part_idx][v]){
					candidate.erase(candidate.find(make_pair(pivot_to_all_dist[part_idx][v],v)));
					pivot_to_all_dist[part_idx][v] = alt;
					candidate.insert(make_pair(alt,v));
				}
			}
		}

		int max_dist = 0;
		for(int k = 0; k < num_nodes; k++){
			if(pivot_to_all_dist[part_idx][k] > max_dist) max_dist = pivot_to_all_dist[part_idx][k];
		}
		cout << "finish partition " << part_idx << " with node: "<< source_node << " with max dist: " << max_dist << endl;
	}

	// output
	// output the distance estimate first
	ofstream partFile(out_dir.append(sample_partition_dist_and_pivot).c_str(),ios::out);
	partFile << num_partition << endl;

	for(int i = 0; i < num_partition; i++){
		for(int j = 0; j < num_partition; j++){
			partFile << indirect_dist[i][j] << " ";
		}
		partFile << endl;
	}

	// output pivot points and distance
	for(int i = 0; i < num_partition; i++){
		partFile << user_max_degree_per_partition[i];
		for(int j = 0; j < num_nodes; j++){
			partFile << " " << pivot_to_all_dist[i][j];
		}
		partFile << endl;
	}
	partFile.close();

}

template <class T1, class T2>
struct dist_smaller{
	bool operator() (const pair<T1,T2>& x, const pair<T1,T2>& y) const {return x.second < y.second;}
};

void Compute2HopNearestNeighbor(vector<vector<pair<int,int> > >& adj_list, string input_dir){
	vector<int> node2hop_nn(adj_list.size(),0);

	#pragma omp parallel for
	for(int node_id = 0; node_id < adj_list.size(); node_id++){
		sort(adj_list[node_id].begin(),adj_list[node_id].end(),dist_smaller<int,int>());

		unordered_set<int> node0_1hop_set;
		node0_1hop_set.insert(node_id);

		// insert all 1 hop neighbour
		for(int j = 0; j < adj_list[node_id].size(); j++){
			node0_1hop_set.insert(adj_list[node_id][j].first);			
		}

		int dist2hop = MAX_DIST;

		for(int j = 0; j < adj_list[node_id].size(); j++){
			int nv = adj_list[node_id][j].first;
			int idx = 0;
			for(idx < adj_list[nv].size(); idx++;){
				if(node0_1hop_set.find(adj_list[nv][idx].first) == node0_1hop_set.end()) break;
			}
			int dist = adj_list[node_id][j].second + adj_list[nv][idx].second;
			if(dist < dist2hop) dist2hop = dist;
		}

		node2hop_nn[node_id] = dist2hop;
	}

	ofstream nn2hopFile(input_dir.append(sample_node2hop_nn),ios::out);
	for(int i = 0; i < adj_list.size(); i++){
		nn2hopFile << node2hop_nn[i] << endl;
	}
	nn2hopFile.close();

}

int main(int argc, char *argv[]){

	if(argc != 5){
		cout << "wrong arguments for " << argv[0] << endl;
		cout << "1.input directory (filtered tweets)" << endl;
		cout << "2.name of partition file index" << endl;
		cout << "3 true partition number" << endl;
		cout << "4 0 for normal partition and 1 for tree partition" << endl;
		exit(1);
	}

	string input_dir(argv[1]);
	int partition_idx = atoi(argv[2]);
	int num_partitions = atoi(argv[3]);
	bool is_tree_partition = atoi(argv[4]);

	vector<vector<pair<int,int> > > adj_list;
	unordered_map<int,int> node_partition_map;
	int true_partitions = LoadGraphPartition(node_partition_map,input_dir,partition_idx);
	if(true_partitions != num_partitions){
		cout << "error in partition loading, true partition not equal to unique partitions" << endl;
		exit(1);
	}
	LoadAdjGraph(adj_list,input_dir);

	if(!is_tree_partition)
	{
		NormalPartition(adj_list,node_partition_map,input_dir,num_partitions);
		Compute2HopNearestNeighbor(adj_list,input_dir);
	}else
	{
		TreePartition(adj_list,node_partition_map,input_dir,num_partitions);
	}
	
}